#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
路径选择器节点 - 文件路径和文件夹路径管理选择显示功能
"""

import os
import json
import folder_paths
from typing import List, Dict, Any, Optional
import comfy.utils

class PathSelectorNode:
    """基础路径选择器节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path_type": (["checkpoints", "loras", "vae", "clip", "clip_vision", "embeddings", "controlnet", "upscale_models", "custom"],),
                "file_filter": ("STRING", {"default": "*", "multiline": False}),
                "include_subdirs": ("BOOLEAN", {"default": False}),
            },
            "optional": {
                "custom_path": ("STRING", {"default": "", "multiline": False}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("selected_path", "file_list", "folder_list")
    FUNCTION = "select_path"
    CATEGORY = "AI训练/path_manager"

    def select_path(self, path_type, file_filter, include_subdirs, custom_path=""):
        """选择路径并返回文件列表"""
        try:
            if path_type == "custom" and custom_path:
                base_path = custom_path
            else:
                base_path = folder_paths.get_folder_paths(path_type)[0] if folder_paths.get_folder_paths(path_type) else ""
            
            if not base_path or not os.path.exists(base_path):
                return ("", "[]", "[]")
            
            file_list = []
            folder_list = []
            
            # 获取文件和文件夹列表
            for item in os.listdir(base_path):
                item_path = os.path.join(base_path, item)
                if os.path.isfile(item_path):
                    if self._matches_filter(item, file_filter):
                        file_list.append(item)
                elif os.path.isdir(item_path) and include_subdirs:
                    folder_list.append(item)
            
            return (base_path, json.dumps(file_list), json.dumps(folder_list))
            
        except Exception as e:
            print(f"路径选择错误: {e}")
            return ("", "[]", "[]")
    
    def _matches_filter(self, filename, file_filter):
        """检查文件是否匹配过滤器"""
        if file_filter == "*":
            return True
        return filename.lower().endswith(file_filter.lower())

class FileBrowserNode:
    """文件浏览器节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "base_path": ("STRING", {"default": "", "multiline": False}),
                "file_extensions": ("STRING", {"default": "*.safetensors,*.ckpt,*.pt", "multiline": False}),
                "recursive": ("BOOLEAN", {"default": False}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("file_paths", "file_names", "folder_structure")
    FUNCTION = "browse_files"
    CATEGORY = "AI训练/path_manager"

    def browse_files(self, base_path, file_extensions, recursive):
        """浏览文件并返回文件列表"""
        try:
            if not base_path or not os.path.exists(base_path):
                return ("[]", "[]", "{}")
            
            extensions = [ext.strip() for ext in file_extensions.split(",")]
            file_paths = []
            file_names = []
            folder_structure = {}
            
            if recursive:
                for root, dirs, files in os.walk(base_path):
                    for file in files:
                        if self._matches_extensions(file, extensions):
                            full_path = os.path.join(root, file)
                            rel_path = os.path.relpath(full_path, base_path)
                            file_paths.append(full_path)
                            file_names.append(rel_path)
                            
                            # 构建文件夹结构
                            folder = os.path.dirname(rel_path)
                            if folder not in folder_structure:
                                folder_structure[folder] = []
                            folder_structure[folder].append(file)
            else:
                for item in os.listdir(base_path):
                    item_path = os.path.join(base_path, item)
                    if os.path.isfile(item_path) and self._matches_extensions(item, extensions):
                        file_paths.append(item_path)
                        file_names.append(item)
                        folder_structure["."] = folder_structure.get(".", []) + [item]
            
            return (json.dumps(file_paths), json.dumps(file_names), json.dumps(folder_structure))
            
        except Exception as e:
            print(f"文件浏览错误: {e}")
            return ("[]", "[]", "{}")
    
    def _matches_extensions(self, filename, extensions):
        """检查文件是否匹配扩展名"""
        for ext in extensions:
            if ext.startswith("*"):
                if filename.lower().endswith(ext[1:].lower()):
                    return True
            elif filename.lower() == ext.lower():
                return True
        return False

class PathValidatorNode:
    """路径验证节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": ("STRING", {"default": "", "multiline": False}),
                "check_type": (["file", "folder", "both"],),
            }
        }
    
    RETURN_TYPES = ("BOOLEAN", "STRING", "STRING")
    RETURN_NAMES = ("is_valid", "error_message", "path_info")
    FUNCTION = "validate_path"
    CATEGORY = "AI训练/path_manager"

    def validate_path(self, path, check_type):
        """验证路径的有效性"""
        try:
            if not path:
                return (False, "路径为空", "{}")
            
            path_info = {
                "exists": os.path.exists(path),
                "is_file": os.path.isfile(path),
                "is_dir": os.path.isdir(path),
                "is_absolute": os.path.isabs(path),
                "size": 0,
                "permissions": {}
            }
            
            if path_info["exists"]:
                if path_info["is_file"]:
                    path_info["size"] = os.path.getsize(path)
                path_info["permissions"] = {
                    "readable": os.access(path, os.R_OK),
                    "writable": os.access(path, os.W_OK),
                    "executable": os.access(path, os.X_OK)
                }
            
            # 根据检查类型判断有效性
            is_valid = False
            error_message = ""
            
            if check_type == "file":
                is_valid = path_info["exists"] and path_info["is_file"]
                if not is_valid:
                    error_message = "路径不是有效的文件"
            elif check_type == "folder":
                is_valid = path_info["exists"] and path_info["is_dir"]
                if not is_valid:
                    error_message = "路径不是有效的文件夹"
            else:  # both
                is_valid = path_info["exists"]
                if not is_valid:
                    error_message = "路径不存在"
            
            return (is_valid, error_message, json.dumps(path_info))
            
        except Exception as e:
            return (False, f"验证错误: {str(e)}", "{}")

class PathMergerNode:
    """路径合并节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "base_path": ("STRING", {"default": "", "multiline": False}),
                "relative_path": ("STRING", {"default": "", "multiline": False}),
                "normalize": ("BOOLEAN", {"default": True}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("merged_path", "path_type")
    FUNCTION = "merge_paths"
    CATEGORY = "AI训练/path_manager"

    def merge_paths(self, base_path, relative_path, normalize):
        """合并路径"""
        try:
            if not base_path:
                return ("", "invalid")
            
            if not relative_path:
                merged_path = base_path
            else:
                merged_path = os.path.join(base_path, relative_path)
            
            if normalize:
                merged_path = os.path.normpath(merged_path)
            
            # 判断路径类型
            if os.path.isfile(merged_path):
                path_type = "file"
            elif os.path.isdir(merged_path):
                path_type = "folder"
            elif os.path.exists(merged_path):
                path_type = "other"
            else:
                path_type = "not_exists"
            
            return (merged_path, path_type)
            
        except Exception as e:
            print(f"路径合并错误: {e}")
            return ("", "error")

# 节点映射
NODE_CLASS_MAPPINGS = {
    "PathSelectorNode": PathSelectorNode,
    "FileBrowserNode": FileBrowserNode,
    "PathValidatorNode": PathValidatorNode,
    "PathMergerNode": PathMergerNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PathSelectorNode": "路径选择器",
    "FileBrowserNode": "文件浏览器",
    "PathValidatorNode": "路径验证器",
    "PathMergerNode": "路径合并器",
} 